package com.chengxiaoxiao.ai_demo.config;

import cn.hutool.core.util.StrUtil;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.listener.ChatModelErrorContext;
import dev.langchain4j.model.chat.listener.ChatModelListener;
import dev.langchain4j.model.chat.listener.ChatModelRequestContext;
import dev.langchain4j.model.chat.listener.ChatModelResponseContext;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.filter.Filter;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.function.Function;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;


/**
 * AI自动配置类
 *
 * @Description:
 * @Author Cheng Xiaoxiao
 * @Date: 2025/3/7 13:54
 */
@Slf4j
@Configuration
@RequiredArgsConstructor
public class AiAutoConfiguration {
    final MySqlChatMemoryStore mySqlChatMemoryStore;

    /**
     * 元数据列：知识库ID
     */
    public static final String META_COLUMN_KNOWLEDGE_BASE_ID = "knowledge_base_id";


    @Bean
    public ChatMemoryProvider chatMemoryProvider() {
        return memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .chatMemoryStore(mySqlChatMemoryStore)
                .maxMessages(10)
                .build();
    }

    @Bean
    ChatModelListener chatModelListener() {
        return new ChatModelListener() {
            @Override
            public void onRequest(ChatModelRequestContext requestContext) {
                log.info("onRequest(): {}", requestContext.chatRequest());
            }

            @Override
            public void onResponse(ChatModelResponseContext responseContext) {
                log.info("onResponse(): {}", responseContext.chatResponse());
            }

            @Override
            public void onError(ChatModelErrorContext errorContext) {
                log.info("onError(): {}", errorContext.error().getMessage());
            }
        };
    }

    @Bean
    public EmbeddingStoreContentRetriever embeddingStoreContentRetriever(ChatLanguageModel chatLanguageModel, EmbeddingModel embeddingModel, EmbeddingStore<TextSegment> embeddingStore, RagProperties ragProperties) {
        Function<Query, Filter> filterByUserId =
                (query) -> {
                    //不为空时，优先使用知识库ID进行查询
                    if (StrUtil.isNotBlank(KnowledgeThreadLocal.getKnowledgeBaseId())) {
                        return metadataKey(META_COLUMN_KNOWLEDGE_BASE_ID).isEqualTo(KnowledgeThreadLocal.getKnowledgeBaseId());
                    }
                    //否则不查询知识库，添加一个一直为false的过滤条件
                    return metadataKey(META_COLUMN_KNOWLEDGE_BASE_ID).isEqualTo("none");
                };

        EmbeddingStoreContentRetriever.EmbeddingStoreContentRetrieverBuilder builder = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .dynamicFilter(filterByUserId);

        if (ragProperties != null) {
            builder
                    .maxResults(ragProperties.getMaxResults())
                    .minScore(ragProperties.getMinScore());
        }

        return builder.build();
    }
}
